# RAG高级---句子窗口检索

import os

from llama_index.readers.web import TrafilaturaWebReader
from llama_index import Document, SimpleDirectoryReader
from llama_index import VectorStoreIndex, StorageContext, load_index_from_storage
from llama_index import load_index_from_storage
from llama_index.readers.web import TrafilaturaWebReader
from llama_index.text_splitter import SentenceSplitter
from llama_index import VectorStoreIndex, ServiceContext
from llama_index.embeddings import resolve_embed_model
from llama_index.node_parser import SentenceWindowNodeParser
from llama_index.indices.postprocessor import MetadataReplacementPostProcessor
from llama_index.indices.postprocessor import SentenceTransformerRerank
from llama_index.llms import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

api_key = "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"
os.environ['OPENAI_API_KEY'] = "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"


# 1、加载文件
documents = SimpleDirectoryReader(
    input_files=["./合同履约-合同结算管理操作手册.pdf"]
).load_data()
document = Document(text="\n\n".join([doc.text for doc in documents]))

# 1、创建句子窗口节点解析器 / 默认设置
node_parser = SentenceWindowNodeParser.from_defaults(
    window_size=3,
    window_metadata_key="window",
    original_text_metadata_key="original_text",
)

# # 测试语句
# text = "你好，很高兴认识你。 已经10点了，可我还不想起床！下雪啦！你的作业完成了吗？"
# # 分词
# nodes = node_parser.get_nodes_from_documents([Document(text=text)])
# print([x.text for x in nodes])

# 2、构建索引
# 2.1 创建OpenAI的llm
llm = OpenAI(model="gpt-3.5-turbo",
             api_key=api_key,
             temperature=0.1)
# 2.2 创建ServiceContext组件
sentence_context = ServiceContext.from_defaults(
    llm=llm,
    embed_model=OpenAIEmbedding(),
    node_parser=node_parser,
)
# 2.3 创建向量数据库
sentence_index = VectorStoreIndex.from_documents(
    [document],
    service_context=sentence_context
)
# 2.4 存储数据至硬盘
sentence_index.storage_context.persist(persist_dir="./sentence_index")


# 3、创建postprocessor组件
# 3.1 创建Replacement组件
postproc = MetadataReplacementPostProcessor(
    target_metadata_key="window"
)
# 3.2 创建rerank组件
# 参考: https://huggingface.co/BAAI/bge-reranker-base
# rerank = SentenceTransformerRerank(
#     top_n=2,
#     model="BAAI/bge-reranker-base"
# )

# 4、创建query engine组件
#创建查询引擎
sentence_window_engine = sentence_index.as_query_engine(
    similarity_top_k=6,
    # node_postprocessors=[postproc,rerank]
    node_postprocessors=[postproc]
)

# 5、提问
window_response = sentence_window_engine.query(
    "发票查询应该怎么做？?"
)
print(window_response)